import torch
import numpy as np
from PIL import Image, ImageOps
import torchvision.transforms as T

class SmartResizeNode:
    """
    ComfyUI custom node for smart image resizing and cropping.
    Automatically handles portrait/landscape images by specifying short and long side values.
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "short_side": ("INT", {
                    "default": 512, 
                    "min": 64, 
                    "max": 2048, 
                    "step": 8
                }),
                "long_side": ("INT", {
                    "default": 768, 
                    "min": 64, 
                    "max": 2048, 
                    "step": 8
                }),
                "method": (["resize", "crop"], {
                    "default": "resize"
                }),
                "crop_position": (["center", "top", "bottom", "left", "right"], {
                    "default": "center"
                }),
            }
        }
    
    RETURN_TYPES = ("IMAGE", "INT", "INT")
    RETURN_NAMES = ("image", "width", "height")
    FUNCTION = "smart_resize"
    CATEGORY = "image/resize"
    
    def smart_resize(self, image, short_side, long_side, method, crop_position):
        # Convert tensor to PIL Image
        batch_size = image.shape[0]
        results = []
        
        for i in range(batch_size):
            # Convert from tensor (B,H,W,C) to PIL Image
            img_tensor = image[i]
            img_array = (img_tensor * 255).clamp(0, 255).byte().cpu().numpy()
            pil_image = Image.fromarray(img_array)
            
            # Get original dimensions
            orig_width, orig_height = pil_image.size
            
            # Determine target dimensions based on orientation
            if orig_width > orig_height:
                # Landscape: width is long side, height is short side
                target_width = long_side
                target_height = short_side
            else:
                # Portrait: height is long side, width is short side
                target_width = short_side
                target_height = long_side
            
            if method == "resize":
                # Resize image maintaining aspect ratio, then pad or crop to exact size
                processed_image = self._smart_resize_with_padding(
                    pil_image, target_width, target_height, crop_position
                )
            else:  # crop
                # Resize and crop to exact dimensions
                processed_image = self._smart_crop(
                    pil_image, target_width, target_height, crop_position
                )
            
            # Convert back to tensor
            img_tensor = torch.from_numpy(np.array(processed_image).astype(np.float32) / 255.0)
            results.append(img_tensor)
        
        # Stack results back into batch tensor
        output_tensor = torch.stack(results, dim=0)
        
        # Get final dimensions from the processed image
        final_height, final_width = output_tensor.shape[1], output_tensor.shape[2]
        
        return (output_tensor, final_width, final_height)
    
    def _smart_resize_with_padding(self, image, target_width, target_height, crop_position):
        """
        Resize image to fit within target dimensions, then pad to exact size
        """
        # Calculate scaling factor to fit image within target dimensions
        scale_w = target_width / image.width
        scale_h = target_height / image.height
        scale = min(scale_w, scale_h)
        
        # Resize image
        new_width = int(image.width * scale)
        new_height = int(image.height * scale)
        resized_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
        
        # Create new image with target dimensions and paste resized image
        result = Image.new('RGB', (target_width, target_height), (0, 0, 0))
        
        # Calculate paste position based on crop_position
        if crop_position == "center":
            paste_x = (target_width - new_width) // 2
            paste_y = (target_height - new_height) // 2
        elif crop_position == "top":
            paste_x = (target_width - new_width) // 2
            paste_y = 0
        elif crop_position == "bottom":
            paste_x = (target_width - new_width) // 2
            paste_y = target_height - new_height
        elif crop_position == "left":
            paste_x = 0
            paste_y = (target_height - new_height) // 2
        elif crop_position == "right":
            paste_x = target_width - new_width
            paste_y = (target_height - new_height) // 2
        
        result.paste(resized_image, (paste_x, paste_y))
        return result
    
    def _smart_crop(self, image, target_width, target_height, crop_position):
        """
        Resize and crop image to exact target dimensions
        """
        # Calculate scaling factor to cover target dimensions
        scale_w = target_width / image.width
        scale_h = target_height / image.height
        scale = max(scale_w, scale_h)
        
        # Resize image
        new_width = int(image.width * scale)
        new_height = int(image.height * scale)
        resized_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
        
        # Calculate crop box based on crop_position
        if crop_position == "center":
            left = (new_width - target_width) // 2
            top = (new_height - target_height) // 2
        elif crop_position == "top":
            left = (new_width - target_width) // 2
            top = 0
        elif crop_position == "bottom":
            left = (new_width - target_width) // 2
            top = new_height - target_height
        elif crop_position == "left":
            left = 0
            top = (new_height - target_height) // 2
        elif crop_position == "right":
            left = new_width - target_width
            top = (new_height - target_height) // 2
        
        right = left + target_width
        bottom = top + target_height
        
        cropped_image = resized_image.crop((left, top, right, bottom))
        return cropped_image

NODE_CLASS_MAPPINGS = {
    "SmartResizeNode": SmartResizeNode
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "SmartResizeNode": "Smart Resize (Short/Long Side)"
}